import sys
import time
import socket
import paramiko


# 添加错误处理、日志记录、更复杂的命令执行逻辑等


class SSHClient:
    def __init__(self, hostname: str,
                 port: int = 22,
                 username: str = 'root',
                 password: str = '123456',
                 timeout: float = 10
                 ):

        self.hostname = hostname
        self.port = port
        self.username = username
        self.password = password
        self.timeout = timeout
        self.client = None
        self.sftp = None

    def conn(self):
        """建立SSH连接"""
        try:
            self.client = paramiko.SSHClient()
            self.client.load_system_host_keys()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            self.client.connect(hostname=self.hostname,
                                port=self.port,
                                username=self.username,
                                password=self.password
                                )
            print(
                f"[INFO] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} SSH connection {self.hostname} established")
        except paramiko.BadHostKeyException as badKey:
            print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} Receive a bad key", badKey)
            sys.exit()
        except paramiko.AuthenticationException as auth:
            print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} Invalid username or password", auth)
            sys.exit()
        except paramiko.SSHException as ssh:
            print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} Establish ssh session error", ssh)
            sys.exit()
        except socket.error as s:
            print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} Connecting has a socket error", s)
            sys.exit()

    def run(self, command: str):
        """执行命令并返回结果"""
        if self.client:
            try:
                stdin, stdout, stderr = self.client.exec_command(command, timeout=self.timeout)
                exit_status = stdout.channel.recv_exit_status()
                if exit_status == 0:
                    return {
                        'stdout': stdout.read().decode('utf-8').strip(),
                        'stderr': stderr.read().decode('utf-8').strip()
                    }

            except socket.timeout:
                raise f"Exec Command {command} timeout"
        else:
            print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} SSH connection not established.")

    def check_port(self, other_port):
        # 检查端口
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(3)
            port_exist = False
            try:
                if other_port:
                    status = sock.connect_ex((self.hostname, other_port))
                    if status == 0:
                        print(f"端口 {other_port} 在 {self.hostname} 上是占用的，请检查")
                        return port_exist
                    else:
                        print(f"端口 {other_port} 在 {self.hostname} 可用")
                        port_exist = True
                else:
                    print("请输入端口")
            except socket.error as e:
                print(print(f"连接失败: {e}"))
        return port_exist

    def sftp_conn(self):
        try:
            t = paramiko.Transport((self.hostname, self.port))
            t.connect(username=self.username, password=self.password)
            self.sftp = paramiko.SFTPClient.from_transport(t)
            print("sftp建立链接")
        except Exception as e:
            print(e)

    def sftp_check_path(self, filepath):
        if self.sftp:
            try:
                self.sftp.stat(filepath)
                print(f"{filepath}目录已存在")
                return False
            except FileNotFoundError:
                print(f"{filepath}目录不存在")
                return True

    def sftp_close(self):
        """断开SFTP连接"""
        if self.sftp:
            self.sftp.close()
            print(f"[INFO] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} SFTP connection closed.")

    def ssh_close(self):
        if self.client:
            self.client.close()
            print(f"[INFO] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} SSH connection closed.")

    def sftp_upload_file(self, local_path, remote_path):
        """上传文件到远程服务器"""
        if self.sftp:
            try:
                print(f"[INFO] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} 正在上传 {local_path} ...")
                self.sftp.put(local_path, remote_path)
                print(
                    f"[INFO] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} 上传文件 {local_path} 到 {self.hostname} 的 {remote_path}.")
            except Exception as e:
                print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} 文件上传失败: {e}")

    def sftp_download(self, remote_path, local_path):
        """从远程服务器下载文件"""
        if self.sftp:
            try:
                self.sftp.get(remote_path, local_path)
                print(
                    f"[INFO] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} 下载 {self.hostname} 的 {remote_path} 到 {local_path}.")
            except Exception as e:
                print(f"[ERROR] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())} 文件下载失败: {e}")

    def read_remote_file(self, filename):
        if self.sftp:
            with self.sftp.open(filename, 'rb') as f:
                return f.read().decode('utf-8')

    def write_remote_file(self, filename, data_str):
        if self.sftp:
            data_str_trans = data_str.replace("\r\n", "\n")
            with self.sftp.open(filename, 'w') as f:
                f.write(data_str_trans)
